/*
 * This file contains balloon driver routines
 *
 * Copyright (c) 2009-2017  Red Hat, Inc.
 *
 * Author(s):
 *  Vadim Rozenfeld <vrozenfe@redhat.com>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met :
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and / or other materials provided with the distribution.
 * 3. Neither the names of the copyright holders nor the names of their contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */
#include "precomp.h"
#include "ntddkex.h"

#if defined(EVENT_TRACING)
#include "Device.tmh"
#endif

#ifdef ALLOC_PRAGMA
#pragma alloc_text(PAGE, BalloonEvtDeviceContextCleanup)
#pragma alloc_text(PAGE, BalloonEvtDevicePrepareHardware)
#pragma alloc_text(PAGE, BalloonEvtDeviceReleaseHardware)
#pragma alloc_text(PAGE, BalloonEvtDeviceD0Exit)
#pragma alloc_text(PAGE, BalloonEvtDeviceD0ExitPreInterruptsDisabled)
#pragma alloc_text(PAGE, BalloonDeviceAdd)
#pragma alloc_text(PAGE, BalloonCloseWorkerThread)
#pragma alloc_text(PAGE, BalloonEvtDeviceSurpriseRemoval)
#pragma alloc_text(PAGE, BalloonCloseMonitorThread)
#endif // ALLOC_PRAGMA

#ifndef BALLOON_INFLATE_IGNORE_LOWMEM
#define LOMEMEVENTNAME L"\\KernelObjects\\LowMemoryCondition"
DECLARE_CONST_UNICODE_STRING(evLowMemString, LOMEMEVENTNAME);
#endif // !BALLOON_INFLATE_IGNORE_LOWMEM

#define VMREADYREGISTRY L"\\Registry\\Machine\\SOFTWARE\\DesktopAppEngine"
DECLARE_CONST_UNICODE_STRING(vmReadyString, VMREADYREGISTRY);


NTSTATUS
BalloonDeviceAdd(
    IN WDFDRIVER  Driver,
    IN PWDFDEVICE_INIT  DeviceInit)
{
    NTSTATUS                     status = STATUS_SUCCESS;
    WDFDEVICE                    device;
    PDEVICE_CONTEXT              devCtx = NULL;
    WDF_INTERRUPT_CONFIG         interruptConfig;
    WDF_OBJECT_ATTRIBUTES        attributes;
    WDF_PNPPOWER_EVENT_CALLBACKS pnpPowerCallbacks;
#ifdef USE_BALLOON_SERVICE
    WDF_FILEOBJECT_CONFIG        fileConfig;
#endif // USE_BALLOON_SERVICE

    UNREFERENCED_PARAMETER(Driver);
    PAGED_CODE();

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "--> %s\n", __FUNCTION__);

    WDF_PNPPOWER_EVENT_CALLBACKS_INIT(&pnpPowerCallbacks);

    pnpPowerCallbacks.EvtDevicePrepareHardware      = BalloonEvtDevicePrepareHardware;
    pnpPowerCallbacks.EvtDeviceReleaseHardware      = BalloonEvtDeviceReleaseHardware;
    pnpPowerCallbacks.EvtDeviceD0Entry              = BalloonEvtDeviceD0Entry;
    pnpPowerCallbacks.EvtDeviceD0Exit               = BalloonEvtDeviceD0Exit;
    pnpPowerCallbacks.EvtDeviceD0ExitPreInterruptsDisabled = BalloonEvtDeviceD0ExitPreInterruptsDisabled;
    pnpPowerCallbacks.EvtDeviceSurpriseRemoval = BalloonEvtDeviceSurpriseRemoval;

    WdfDeviceInitSetPnpPowerEventCallbacks(DeviceInit, &pnpPowerCallbacks);

    WDF_OBJECT_ATTRIBUTES_INIT_CONTEXT_TYPE(&attributes, DEVICE_CONTEXT);
    attributes.EvtCleanupCallback = BalloonEvtDeviceContextCleanup;

    /* The driver initializes all the queues under lock of
     * WDF object of the device. If we use default execution
     * level, this lock is spinlock and common blocks required
     * for queues can't be allocated on DISPATCH. So, we change
     * the execution level to PASSIVE -> the lock is fast mutex
     */
    attributes.ExecutionLevel = WdfExecutionLevelPassive;

#ifdef USE_BALLOON_SERVICE
    attributes.SynchronizationScope = WdfSynchronizationScopeDevice;

    WDF_FILEOBJECT_CONFIG_INIT(
                            &fileConfig,
                            WDF_NO_EVENT_CALLBACK,
                            BalloonEvtFileClose,
                            WDF_NO_EVENT_CALLBACK
                            );

    WdfDeviceInitSetFileObjectConfig(DeviceInit,
                            &fileConfig,
                            WDF_NO_OBJECT_ATTRIBUTES);
#endif // USE_BALLOON_SERVICE

    status = WdfDeviceCreate(&DeviceInit, &attributes, &device);
    if(!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
           "WdfDeviceCreate failed with status 0x%08x\n", status);
        return status;
    }

    devCtx = GetDeviceContext(device);

    WDF_INTERRUPT_CONFIG_INIT(&interruptConfig,
                            BalloonInterruptIsr,
                            BalloonInterruptDpc);

    interruptConfig.EvtInterruptEnable  = BalloonInterruptEnable;
    interruptConfig.EvtInterruptDisable = BalloonInterruptDisable;

#ifndef USE_BALLOON_SERVICE
    // Serialize the DPC routine with queue operations to prevent races
    // around PendingWriteRequest and HandleWriteRequest.
    interruptConfig.AutomaticSerialization = TRUE;
#endif // !USE_BALLOON_SERVICE

    status = WdfInterruptCreate(device,
                            &interruptConfig,
                            WDF_NO_OBJECT_ATTRIBUTES,
                            &devCtx->WdfInterrupt);
    if (!NT_SUCCESS (status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
           "WdfInterruptCreate failed: 0x%08x\n", status);
        return status;
    }

    status = WdfDeviceCreateDeviceInterface(device, &GUID_DEVINTERFACE_BALLOON, NULL);
    if(!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
           "WdfDeviceCreateDeviceInterface failed with status 0x%08x\n", status);
        return status;
    }

    devCtx->SurpriseRemoval = FALSE;
    devCtx->bShutDown = FALSE;
    devCtx->num_pages = 0;
    devCtx->PageListHead.Next = NULL;
    ExInitializeNPagedLookasideList(
                      &devCtx->LookAsideList,
                      NULL,
                      NULL,
#if !defined(NTDDI_WIN8) || (NTDDI_VERSION < NTDDI_WIN8)
                      0,
#else
                      POOL_NX_ALLOCATION,
#endif
                      sizeof(PAGE_LIST_ENTRY),
                      BALLOON_MGMT_POOL_TAG,
                      0
                      );
    devCtx->bListInitialized = TRUE;
    devCtx->pfns_table = NULL;
    devCtx->MemStats = NULL;

    KeInitializeEvent(&devCtx->HostAckEvent,
                      SynchronizationEvent,
                      FALSE
                      );

    WDF_OBJECT_ATTRIBUTES_INIT(&attributes);
    attributes.ParentObject = device;

    status = WdfSpinLockCreate(
        &attributes,
        &devCtx->StatQueueLock
        );
    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP, "WdfSpinLockCreate failed 0x%x\n", status);
        return status;
    }

    status = WdfSpinLockCreate(
        &attributes,
        &devCtx->InfDefQueueLock
        );
    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP, "WdfSpinLockCreate failed 0x%x\n", status);
        return status;
    }

#ifdef USE_BALLOON_SERVICE
    status = BalloonQueueInitialize(device);
    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
           "BalloonQueueInitialize failed with status 0x%08x\n", status);
        return status;
    }
#else // USE_BALLOON_SERVICE
    status = StatInitializeWorkItem(device);
    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
           "StatInitializeWorkItem failed with status 0x%08x\n", status);
        return status;
    }
#endif // USE_BALLOON_SERVICE

    KeInitializeEvent(&devCtx->WakeUpThread,
                      SynchronizationEvent,
                      FALSE
                      );

    // Auto-Balloon structure init
    status = WdfSpinLockCreate(
        &attributes,
        &devCtx->MessageQueueLock);
    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP, "WdfSpinLockCreate failed 0x%x\n", status);
        return status;
    }
    KeInitializeEvent(&devCtx->MessageAckEvent,
                      SynchronizationEvent,
                      FALSE);
    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "<-- %s\n", __FUNCTION__);
    return status;
}

VOID
BalloonEvtDeviceContextCleanup(
    IN WDFOBJECT  Device
    )
{
    PDEVICE_CONTEXT     devCtx = GetDeviceContext((WDFDEVICE)Device);

    PAGED_CODE();

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "--> %s\n", __FUNCTION__);

    if(devCtx->bListInitialized)
    {
        ExDeleteNPagedLookasideList(&devCtx->LookAsideList);
        devCtx->bListInitialized = FALSE;
    }

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "<-- %s\n", __FUNCTION__);
}

NTSTATUS
BalloonEvtDevicePrepareHardware(
    IN WDFDEVICE    Device,
    IN WDFCMRESLIST ResourceList,
    IN WDFCMRESLIST ResourceListTranslated
    )
{
    NTSTATUS            status         = STATUS_SUCCESS;
    PDEVICE_CONTEXT     devCtx = NULL;

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "--> %s\n", __FUNCTION__);

    UNREFERENCED_PARAMETER(ResourceList);

    PAGED_CODE();

    devCtx = GetDeviceContext(Device);

    status = VirtIOWdfInitialize(&devCtx->VDevice, Device, ResourceListTranslated, NULL, BALLOON_MGMT_POOL_TAG);
    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_POWER, "VirtIOWdfInitialize failed with %x\n", status);
        return status;
    }

    if (NT_SUCCESS(status))
    {
        devCtx->MemStats = (PBALLOON_STAT)VirtIOWdfDeviceAllocDmaMemory(&devCtx->VDevice.VIODevice, PAGE_SIZE, BALLOON_MGMT_POOL_TAG);
    }

    if (devCtx->MemStats)
    {
        RtlFillMemory(devCtx->MemStats, sizeof(BALLOON_STAT) * VIRTIO_BALLOON_S_NR, -1);
    }
    else
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_POWER, "Failed to allocate MemStats block\n");
        status = STATUS_INSUFFICIENT_RESOURCES;
    }

    /* use BALLOON_MGMT_POOL_TAG also for tagging common memory blocks */
    if (NT_SUCCESS(status))
    {
        devCtx->pfns_table = (PPFN_NUMBER)VirtIOWdfDeviceAllocDmaMemory(&devCtx->VDevice.VIODevice, PAGE_SIZE, BALLOON_MGMT_POOL_TAG);
    }

    if (devCtx->pfns_table == NULL)
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP, "Failed to allocate PFNS_TABLE block\n");
        status = STATUS_INSUFFICIENT_RESOURCES;
        return status;
    }

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "<-- %s\n", __FUNCTION__);
    return status;
}

NTSTATUS
BalloonEvtDeviceReleaseHardware (
    IN WDFDEVICE      Device,
    IN WDFCMRESLIST   ResourcesTranslated
    )
{
    PDEVICE_CONTEXT     devCtx = NULL;

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "--> %s\n", __FUNCTION__);

    UNREFERENCED_PARAMETER(ResourcesTranslated);

    PAGED_CODE();

    WdfObjectAcquireLock(Device);

    devCtx = GetDeviceContext(Device);

    VirtIOWdfDeviceFreeDmaMemoryByTag(&devCtx->VDevice.VIODevice, BALLOON_MGMT_POOL_TAG);
    devCtx->MemStats = NULL;
    devCtx->pfns_table = NULL;

    WdfObjectReleaseLock(Device);

    VirtIOWdfShutdown(&devCtx->VDevice);

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "<-- %s\n", __FUNCTION__);
    return STATUS_SUCCESS;
}

NTSTATUS
BalloonCreateWorkerThread(
    IN WDFDEVICE  Device
    )
{
    PDEVICE_CONTEXT     devCtx = GetDeviceContext(Device);
    NTSTATUS            status = STATUS_SUCCESS;
    HANDLE              hThread = 0;
    OBJECT_ATTRIBUTES   oa;

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "--> %s\n", __FUNCTION__);
    devCtx->bShutDown = FALSE;

    if(NULL == devCtx->Thread)
    {
        InitializeObjectAttributes(&oa, NULL, 
            OBJ_KERNEL_HANDLE, NULL, NULL);

        status = PsCreateSystemThread(&hThread, THREAD_ALL_ACCESS, &oa, NULL, NULL,
                                          BalloonRoutine, Device);

        if(!NT_SUCCESS(status))
        {
            TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
              "failed to create worker thread status 0x%08x\n", status);
            return status;
        }

        ObReferenceObjectByHandle(hThread, THREAD_ALL_ACCESS, NULL,
                                  KernelMode, (PVOID*)&devCtx->Thread, NULL);
        KeSetPriorityThread(devCtx->Thread, LOW_REALTIME_PRIORITY);

        ZwClose(hThread);
    }

    KeSetEvent(&devCtx->WakeUpThread, EVENT_INCREMENT, FALSE);

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "<-- %s\n", __FUNCTION__);
    return status;
}

NTSTATUS
BalloonCloseWorkerThread(
    IN WDFDEVICE  Device
    )
{
    PDEVICE_CONTEXT     devCtx = GetDeviceContext(Device);
    NTSTATUS            status = STATUS_SUCCESS;

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "--> %s\n", __FUNCTION__);

    PAGED_CODE();

    if(NULL != devCtx->Thread)
    {
        devCtx->bShutDown = TRUE;
        KeSetEvent(&devCtx->WakeUpThread, EVENT_INCREMENT, FALSE);
        status = KeWaitForSingleObject(devCtx->Thread, Executive, KernelMode, FALSE, NULL);
        if(!NT_SUCCESS(status))
        {
            TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
              "KeWaitForSingleObject didn't succeed status 0x%08x\n", status);
        }
        ObDereferenceObject(devCtx->Thread);
        devCtx->Thread = NULL;
    }

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "<-- %s\n", __FUNCTION__);
    return status;
}


NTSTATUS
BalloonEvtDeviceD0Entry(
    IN  WDFDEVICE Device,
    IN  WDF_POWER_DEVICE_STATE PreviousState
    )
{
    NTSTATUS            status = STATUS_SUCCESS;
    PDEVICE_CONTEXT devCtx = GetDeviceContext(Device);

    UNREFERENCED_PARAMETER(PreviousState);
    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "--> %s\n", __FUNCTION__);

    status = BalloonInit(Device);
    if(!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
           "BalloonInit failed with status 0x%08x\n", status);
        BalloonTerm(Device);
        return status;
    }

    status = BalloonCreateWorkerThread(Device);
    if(!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
           "BalloonCreateWorkerThread failed with status 0x%08x\n", status);
        BalloonTerm(Device);
        return status;
    } 

#ifndef BALLOON_INFLATE_IGNORE_LOWMEM
    devCtx->evLowMem = IoCreateNotificationEvent(
        (PUNICODE_STRING)&evLowMemString, &devCtx->hLowMem);
#endif // !BALLOON_INFLATE_IGNORE_LOWMEM

    // Create Auto-Balloon monitor thread
    if (devCtx->AutoBalloonContext != NULL)
    {
        status = BalloonCreateMonitorThread(Device);
        if (!NT_SUCCESS(status))
        {
            TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
                "BalloonCreateMonitorThread failed with status 0x%08x\n", status);
            BalloonTerm(Device);
            BalloonCloseWorkerThread(Device);
        }
    }
    return status;
}

NTSTATUS
BalloonEvtDeviceD0Exit(
    IN  WDFDEVICE Device,
    IN  WDF_POWER_DEVICE_STATE TargetState
    )
{
    PDEVICE_CONTEXT devCtx = GetDeviceContext(Device);

    UNREFERENCED_PARAMETER(TargetState);

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "<--> %s\n", __FUNCTION__);

    PAGED_CODE();

#ifndef BALLOON_INFLATE_IGNORE_LOWMEM
    if (devCtx->evLowMem)
    {
        ZwClose(devCtx->hLowMem);
        devCtx->evLowMem = NULL;
    }
#endif // !BALLOON_INFLATE_IGNORE_LOWMEM

#ifndef USE_BALLOON_SERVICE
   /*
    * interrupts were already disabled (between BalloonEvtDeviceD0ExitPreInterruptsDisabled and this call)
    * we should flush StatWorkItem before calling BalloonTerm which will delete virtio queues
    */
    WdfWorkItemFlush(devCtx->StatWorkItem);
#endif // !USE_BALLOON_SERVICE

    BalloonTerm(Device);

    return STATUS_SUCCESS;
}

NTSTATUS
BalloonEvtDeviceD0ExitPreInterruptsDisabled(
    IN  WDFDEVICE Device,
    IN  WDF_POWER_DEVICE_STATE TargetState
    )
{
    PDEVICE_CONTEXT       devCtx = GetDeviceContext(Device);
    POWER_ACTION    powerAction = WdfDeviceGetSystemPowerAction(Device);
    BOOLEAN sysPowerS4Final = FALSE;

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "<--> %s\n", __FUNCTION__);

    PAGED_CODE();

    // https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wdm/ne-wdm-power_action
    switch (powerAction) {
        case PowerActionHibernate:
        case PowerActionShutdown:
        case PowerActionShutdownReset:
        case PowerActionShutdownOff:
            TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "--> %s\n", __FUNCTION__);
            sysPowerS4Final = TRUE;
            break;
        default:
            break;
    }

    BalloonCloseWorkerThread(Device);
    // Destroy Auto-Balloon monitor thread
    BalloonCloseMonitorThread(Device);

    if (TargetState == WdfPowerDeviceD3Final && sysPowerS4Final == FALSE)
    {
       while(devCtx->num_pages)
       {
          BalloonLeak(Device, devCtx->num_pages);
       }

       BalloonSetSize(Device, devCtx->num_pages);
    }
    return STATUS_SUCCESS;
}

VOID
BalloonEvtDeviceSurpriseRemoval(IN WDFDEVICE Device)
{
    PDEVICE_CONTEXT devCtx = GetDeviceContext(Device);
    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "<--> %s\n", __FUNCTION__);

    PAGED_CODE();

    devCtx->SurpriseRemoval = TRUE;
    KeSetEvent(&devCtx->HostAckEvent, EVENT_INCREMENT, FALSE);
}

BOOLEAN
BalloonInterruptIsr(
    IN WDFINTERRUPT WdfInterrupt,
    IN ULONG        MessageID
    )
{
    PDEVICE_CONTEXT     devCtx = NULL;
    WDFDEVICE           Device;

    UNREFERENCED_PARAMETER( MessageID );

    Device = WdfInterruptGetDevice(WdfInterrupt);
    devCtx = GetDeviceContext(Device);

    if (VirtIOWdfGetISRStatus(&devCtx->VDevice) > 0)
    {
        TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INTERRUPT, "--> %s\n", __FUNCTION__);
        WdfInterruptQueueDpcForIsr( WdfInterrupt );
        return TRUE;
    }
    else
    {
        TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INTERRUPT, "--> %s No Isr indicated\n", __FUNCTION__);
    }
    return FALSE;
}

VOID
BalloonInterruptDpc(
    IN WDFINTERRUPT WdfInterrupt,
    IN WDFOBJECT    WdfDevice
    )
{
    unsigned int          len;
    PDEVICE_CONTEXT       devCtx = GetDeviceContext(WdfDevice);
    PVOID                 buffer;

    BOOLEAN               bHostAck = FALSE;
    UNREFERENCED_PARAMETER( WdfInterrupt );

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_DPC, "--> %s\n", __FUNCTION__);

    WdfSpinLockAcquire(devCtx->InfDefQueueLock);
    if (virtqueue_get_buf(devCtx->InfVirtQueue, &len))
    {
        bHostAck = TRUE;
    }
    if (virtqueue_get_buf(devCtx->DefVirtQueue, &len))
    {
        bHostAck = TRUE;
    }
    WdfSpinLockRelease(devCtx->InfDefQueueLock);

    if(bHostAck)
    {
        KeSetEvent(&devCtx->HostAckEvent, EVENT_INCREMENT, FALSE);
    }

    if (devCtx->MessageVirtQueue)
    {
        WdfSpinLockAcquire(devCtx->MessageQueueLock);
        buffer = virtqueue_get_buf(devCtx->MessageVirtQueue, &len);
        WdfSpinLockRelease(devCtx->MessageQueueLock);
        if (buffer)
        {
            KeSetEvent(&devCtx->MessageAckEvent, EVENT_INCREMENT, FALSE);
        }
    }

    if (devCtx->StatVirtQueue)
    {
        WdfSpinLockAcquire(devCtx->StatQueueLock);
        buffer = virtqueue_get_buf(devCtx->StatVirtQueue, &len);
        WdfSpinLockRelease(devCtx->StatQueueLock);

        if (buffer)
        {
#ifdef USE_BALLOON_SERVICE
            WDFREQUEST request = devCtx->PendingWriteRequest;

            devCtx->HandleWriteRequest = TRUE;

            if ((request != NULL) &&
                (WdfRequestUnmarkCancelable(request) != STATUS_CANCELLED))
            {
                NTSTATUS status;
                PVOID buffer;
                size_t length = 0;

                devCtx->PendingWriteRequest = NULL;

                status = WdfRequestRetrieveInputBuffer(request, 0, &buffer, &length);
                if (!NT_SUCCESS(status))
                {
                    length = 0;
                }
                WdfRequestCompleteWithInformation(request, status, length);
            }
#else // USE_BALLOON_SERVICE
            /*
             * According to MSDN 'Using Framework Work Items' article:
             * Create one or more work items that your driver requeues as necessary.
             * Subsequently, each time that the driver's EvtInterruptDpc callback
             * function is called it must determine if the EvtWorkItem callback
             * function has run. If the EvtWorkItem callback function has not run,
             * the EvtInterruptDpc callback function does not call WdfWorkItemEnqueue,
             * because the work item is still queued.
             * A few drivers might need to call WdfWorkItemFlush to flush their work
             * items from the work-item queue.
             *
             * For each dpc (i.e. interrupt) we'll push stats exactly that many times.
             */
            if (1==InterlockedIncrement(&devCtx->WorkCount))
            {
                WdfWorkItemEnqueue(devCtx->StatWorkItem);
            }
#endif // USE_BALLOON_SERVICE
        }
    }

    if(devCtx->Thread)
    {
       KeSetEvent(&devCtx->WakeUpThread, EVENT_INCREMENT, FALSE);
    }
}

NTSTATUS
BalloonInterruptEnable(
    IN WDFINTERRUPT WdfInterrupt,
    IN WDFDEVICE    WdfDevice
    )
{
    PDEVICE_CONTEXT     devCtx = NULL;

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_PNP, "--> %s\n", __FUNCTION__);

    devCtx = GetDeviceContext(WdfDevice);
    EnableInterrupt(WdfInterrupt, devCtx);
    BalloonInterruptIsr(WdfInterrupt, 0);

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_PNP, "<-- %s\n", __FUNCTION__);
    return STATUS_SUCCESS;
}

NTSTATUS
BalloonInterruptDisable(
    IN WDFINTERRUPT WdfInterrupt,
    IN WDFDEVICE    WdfDevice
    )
{
    PDEVICE_CONTEXT     devCtx = NULL;
    UNREFERENCED_PARAMETER( WdfInterrupt );

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_PNP, "--> %s\n", __FUNCTION__);

    devCtx = GetDeviceContext(WdfDevice);
    DisableInterrupt(devCtx);

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_PNP, "<-- %s\n", __FUNCTION__);
    return STATUS_SUCCESS;
}

#ifdef USE_BALLOON_SERVICE
VOID
BalloonEvtFileClose(
    IN WDFFILEOBJECT    FileObject
    )
{
    WDFDEVICE Device = WdfFileObjectGetDevice(FileObject);
    PDEVICE_CONTEXT devCtx = GetDeviceContext(Device);

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "<-> %s\n", __FUNCTION__);

    // synchronize with the device to make sure it doesn't exit D0 from underneath us
    WdfObjectAcquireLock(Device);

    if (devCtx->MemStats)
    {
        RtlFillMemory(devCtx->MemStats, sizeof(BALLOON_STAT) * VIRTIO_BALLOON_S_NR, -1);
    }

    if (devCtx->StatVirtQueue && devCtx->MemStats)
    {
        BalloonMemStats(Device);
    }

    WdfObjectReleaseLock(Device);
}
#endif // USE_BALLOON_SERVICE

VOID
BalloonSetSize(
    IN WDFOBJECT WdfDevice,
    IN size_t    num
    )
{
    PDEVICE_CONTEXT       devCtx = GetDeviceContext(WdfDevice);
    u32 actual = (u32)num;
    VirtIOWdfDeviceSet(&devCtx->VDevice, FIELD_OFFSET(VIRTIO_BALLOON_CONFIG, actual), &actual, sizeof(actual));
}

LONGLONG
BalloonGetSize(
    IN WDFOBJECT WdfDevice
    )
{
    PDEVICE_CONTEXT       devCtx = GetDeviceContext(WdfDevice);

    u32 v;
    VirtIOWdfDeviceGet(&devCtx->VDevice, FIELD_OFFSET(VIRTIO_BALLOON_CONFIG, num_pages), &v, sizeof(v));
    return (LONGLONG)v - devCtx->num_pages;
}

VOID
BalloonRoutine(
    IN PVOID pContext
    )
{
    WDFOBJECT Device  =  (WDFOBJECT)pContext;
    PDEVICE_CONTEXT                 devCtx = GetDeviceContext(Device);

    NTSTATUS            status = STATUS_SUCCESS;
    LONGLONG            diff;

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "Balloon thread started....\n");

    for (;;)
    {
        status = KeWaitForSingleObject(&devCtx->WakeUpThread, Executive,
                                       KernelMode, FALSE, NULL);
        if(STATUS_WAIT_0 == status)
        {
            if(devCtx->bShutDown)
            {
                TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "Exiting Thread!\n");
                break;
            }
            else
            {
                diff = BalloonGetSize(Device);
                if (diff > 0)
                {
                    BalloonFill(Device, (size_t)(diff));
                }
                else if (diff < 0)
                {
                    BalloonLeak(Device, (size_t)(-diff));
                }
                BalloonSetSize(Device, devCtx->num_pages);
            }
        }
    }
    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_PNP, "Thread about to exit...\n");

    PsTerminateSystemThread(STATUS_SUCCESS);
}

/* AutoBalloon Feature */
typedef struct _AUTO_BALLOON_STATS {
    ULONGLONG IoReadTransferCount;
    ULONG CommittedPages;
    ULONG AvailablePages; // Sum (Standby List + Free List + Zero List)
    ULONG PageReadCount;
    ULONG ActualPages;  // Number of pages actually occupied by balloon driver
    LARGE_INTEGER LocalTime;
} AUTO_BALLOON_STATS, *PAUTO_BALLOON_STATS;

typedef struct _AUTO_BALLOON_PRED {
    ULONGLONG LastIoReadTransferCount;
    ULONG InUsedPages;
    ULONG PageReadCount;
    ULONG CommittedPages;
    LARGE_INTEGER LocalTime;
} AUTO_BALLOON_PRED, *PAUTO_BALLOON_PRED;

static BOOLEAN GatherAutoBalloonStats(OUT AUTO_BALLOON_STATS *stats)
{
    NTSTATUS ntStatus;
    SYSTEM_PERFORMANCE_INFORMATION perfInfo;
    ULONG outLen = 0;

    RtlZeroMemory(&perfInfo, sizeof(perfInfo));

    ntStatus = ZwQuerySystemInformation(SystemPerformanceInformation, &perfInfo, sizeof(perfInfo), &outLen);
    if (!NT_SUCCESS(ntStatus)) {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_UVP_DEBUG,
            "TraceKernelStats (SystemPerformanceInformation) call Error\n");
        return FALSE;
    }
    if (outLen != sizeof(perfInfo)) {
        TraceEvents(TRACE_LEVEL_WARNING, DBG_UVP_DEBUG,
            "TraceKernelStats (SystemPerformanceInformation) expected outLen=0x%08x returned with 0x%0x\n",
            sizeof(perfInfo), outLen);
        return FALSE;
    }

    stats->IoReadTransferCount = perfInfo.IoReadTransferCount.QuadPart;
    stats->CommittedPages = perfInfo.CommittedPages - stats->ActualPages;
    stats->AvailablePages = perfInfo.AvailablePages + stats->ActualPages;
    stats->PageReadCount = perfInfo.PageReadCount;

    KeQuerySystemTime(&stats->LocalTime);
    return TRUE;
}

static VOID UpdatePred(
    IN AUTO_BALLOON_PRED *pred,
    IN AUTO_BALLOON_STATS *stats,
    ULONG newInUsed,
    ULONG newCommitted)
{
    pred->LastIoReadTransferCount = stats->IoReadTransferCount;
    pred->InUsedPages = newInUsed;
    pred->CommittedPages = newCommitted;
    pred->PageReadCount = stats->PageReadCount;
    pred->LocalTime = stats->LocalTime;
}

static ULONG AutoBalloonPredict(
    IN WDFDEVICE Device,
    IN AUTO_BALLOON_PRED *pred,
    IN AUTO_BALLOON_STATS *stats
    )
{
    PDEVICE_CONTEXT devCtx = GetDeviceContext(Device);
    AUTO_BALLOON_CONTEXT *autoBalloon = devCtx->AutoBalloonContext;
    ULONG inUsed = autoBalloon->RamPages - stats->AvailablePages;
    ULONG ioReadDelta = (ULONG)(stats->IoReadTransferCount - pred->LastIoReadTransferCount);
    ULONG newInUsed = 0, newAvailable = 0, actualPages = 0, newCommitted = 0, newBalloonPages = 0;
    ULONG masterInUsed = max(pred->InUsedPages, inUsed);
    ULONG slaverInUsed = min(pred->InUsedPages, inUsed);
    ULONG masterCommitted = max(pred->CommittedPages, stats->CommittedPages);
    ULONG slaverCommitted = min(stats->CommittedPages, pred->CommittedPages);
    ULONG balloonPages = devCtx->num_pages;
    LONGLONG timeInterval = stats->LocalTime.QuadPart - pred->LocalTime.QuadPart;

    // 1. Predict InUsed pages
    if (stats->CommittedPages > pred->CommittedPages) {
        // 1.1 increase next = (weight + 1 / weight - 1) * current
        /*
         * Try to deflate next moment, predict memory required by the system
         * based on the maximum values of periodic data and
         * current system memory status to deflate more memory and ensure system performance.
         */
        newInUsed = masterInUsed * MEMORY_INUSE_WEIGHT / WEIGHT_SCALE +
            masterInUsed * WEIGHT_SCALE / MEMORY_INUSE_WEIGHT - masterInUsed;
        newCommitted = masterCommitted * MEMORY_INUSE_WEIGHT / WEIGHT_SCALE +
            masterCommitted * WEIGHT_SCALE / MEMORY_INUSE_WEIGHT - masterCommitted;
    } else {
        // 1.2 descent next = weight * master + (1 - weight) * slaver
        /*
         * Try to inflate next moment, predict memory required by the system
         * based on the periodic value and current system memory status.
         */
        newInUsed = (MEMORY_INUSE_WEIGHT * masterInUsed +
            (WEIGHT_SCALE - MEMORY_INUSE_WEIGHT) * slaverInUsed) / WEIGHT_SCALE;
        newCommitted = (MEMORY_INUSE_WEIGHT * masterCommitted +
            (WEIGHT_SCALE - MEMORY_INUSE_WEIGHT) * slaverCommitted) / WEIGHT_SCALE;
    }
    // 1.3 check InUsed
    if (newInUsed < newCommitted / ARITHMETIC_MEAN) {
        newInUsed = newCommitted / ARITHMETIC_MEAN;
    }

    if (newInUsed < inUsed) {
        newInUsed = inUsed;
    }

    // 2.1 Increase InUsed pages when swap in
    if ((stats->PageReadCount - pred->PageReadCount) >
        (HARD_PAGE_FAULT_THRESHOLD * timeInterval / SECONDS_FROM_100NANOSECONDS)) {
        newInUsed += stats->PageReadCount - pred->PageReadCount;
    }
    // 2.2 Increase InUsed pages when IoReadTransferCount increases significantly.
    if (ioReadDelta > IO_READ_TRANSFER_THRESHOLD * timeInterval / SECONDS_FROM_100NANOSECONDS) {
        newInUsed += ioReadDelta * IO_READ_TRANSFER_RATIO;
    }
    TraceEvents(TRACE_LEVEL_ERROR, DBG_UVP_DEBUG,
            "IoReadTransferCount: threshold is %lld, ioReadDelta %lld.\n", IO_READ_TRANSFER_THRESHOLD, ioReadDelta);

    // 3. Available pages
    actualPages = newInUsed + newInUsed * autoBalloon->BufferWeight / WEIGHT_SCALE;
    // 4. The rest of memory is balloon pages
    newBalloonPages = autoBalloon->RamPages > actualPages ? (autoBalloon->RamPages - actualPages) : 0;

    // 5. Update predict info
    UpdatePred(pred, stats, newInUsed, newCommitted);

    if ((balloonPages != 0) && (newBalloonPages != 0)) {
        // Predict balloon size based on EWMA.
        newBalloonPages = ((CALCULATE_BALLOON_SIZE_WEIGHT * newBalloonPages) +
            (WEIGHT_SCALE - CALCULATE_BALLOON_SIZE_WEIGHT) * balloonPages) / WEIGHT_SCALE;
    }
    newBalloonPages = min(devCtx->AutoBalloonContext->InflateMax, newBalloonPages);
    return newBalloonPages;
}

/*
 * Wake up the main thread in edge-triggered mode.
 * If IsReady changes from 0 to 1, the memory optimization feature is enabled.
 * If IsReady changes from 1 to 0, the memory optimization feature is disabled.
 * Other values are ignored.
 * If we fail to query the registry, we will disable the memory optimization feature.
 */
static BOOLEAN QueryRegistry(IN WDFDEVICE  Device)
{
    PDEVICE_CONTEXT devCtx = GetDeviceContext(Device);
    NTSTATUS ntStatus;
    UNICODE_STRING regKey;
    OBJECT_ATTRIBUTES ObjectAttributes;
    HANDLE hRegistry = NULL;
    ULONG outLen;
    UCHAR regValue[sizeof(KEY_VALUE_PARTIAL_INFORMATION) + sizeof(DWORD)];
    UINT32 newIsReady = 0;

    RtlInitUnicodeString(&regKey, L"IsReady");
    InitializeObjectAttributes(&ObjectAttributes,
        (PUNICODE_STRING)&vmReadyString,
        OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE,
        NULL,
        NULL);
    ntStatus = ZwOpenKey(&hRegistry, KEY_ALL_ACCESS, &ObjectAttributes);
    if (!NT_SUCCESS(ntStatus))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_UVP_DEBUG,
            "Open register key failed 0x%08x \n", ntStatus);
        return FALSE;
    }

    ntStatus = ZwQueryValueKey(hRegistry, &regKey,
        KeyValuePartialInformation, regValue, sizeof(regValue), &outLen);
    if (!NT_SUCCESS(ntStatus))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_UVP_DEBUG,
            "Query register key failed 0x%08x \n", ntStatus);
        ZwClose(hRegistry);
        return FALSE;
    }
    ZwClose(hRegistry);
    if (memcpy_s(&newIsReady, sizeof(newIsReady), ((PKEY_VALUE_PARTIAL_INFORMATION)regValue)->Data,
        ((PKEY_VALUE_PARTIAL_INFORMATION)regValue)->DataLength) != 0)
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_UVP_DEBUG,
            "Copy registry failed\n");
        return FALSE;
    }
    if (newIsReady < VM_STATE_MAX && devCtx->AutoBalloonContext->IsReady != newIsReady)
    {
        devCtx->AutoBalloonContext->IsReady = newIsReady;
        return TRUE;
    }
    return FALSE;
}

static VOID BalloonMemoryInit(
    IN WDFDEVICE  Device,
    OUT AUTO_BALLOON_STATS* currStats,
    OUT AUTO_BALLOON_PRED* inflatePred,
    OUT AUTO_BALLOON_PRED* deflatePred
    )
{
    NTSTATUS ntStatus;
    PDEVICE_CONTEXT devCtx = GetDeviceContext(Device);

    inflatePred->CommittedPages = currStats->CommittedPages;
    inflatePred->InUsedPages = devCtx->AutoBalloonContext->RamPages - currStats->AvailablePages;
    inflatePred->PageReadCount = currStats->PageReadCount;
    inflatePred->LocalTime.QuadPart = currStats->LocalTime.QuadPart;

    deflatePred->CommittedPages = currStats->CommittedPages;
    deflatePred->InUsedPages = devCtx->AutoBalloonContext->RamPages - currStats->AvailablePages;
    deflatePred->PageReadCount = currStats->PageReadCount;
    deflatePred->LocalTime.QuadPart = currStats->LocalTime.QuadPart;

    devCtx->AutoBalloonContext->InflateMax =
        (devCtx->AutoBalloonContext->InflateWeight * devCtx->AutoBalloonContext->RamPages) / WEIGHT_SCALE;
    devCtx->AutoBalloonContext->bCtxInitialized = TRUE;
}

static VOID TryInflate(IN WDFDEVICE  Device)
{
    PDEVICE_CONTEXT devCtx = GetDeviceContext(Device);
    if (devCtx->num_pages != devCtx->AutoBalloonContext->InflateMax) {
        devCtx->AutoBalloonContext->Actual->val = devCtx->AutoBalloonContext->InflateMax;
        BalloonTellMessage(Device);
    }
}

// Auto-Balloon monitor
static VOID AutoBalloonWorker(IN PVOID pContext)
{
    WDFOBJECT Device = (WDFOBJECT)pContext;
    PDEVICE_CONTEXT devCtx = GetDeviceContext(Device);
    NTSTATUS ntStatus;
    LARGE_INTEGER delay;
    ULONG balloonPages;
    ULONG AvailablePagesPeak = devCtx->AutoBalloonContext->RamPages;
    ULONG CommittedPagesPeak = 0;
    AUTO_BALLOON_STATS currStats = { 0 };
    AUTO_BALLOON_PRED inflatePred = { 0 };
    AUTO_BALLOON_PRED deflatePred = { 0 };

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_UVP_DEBUG, "Auto-Balloon thread started....");

    delay.QuadPart = WDF_REL_TIMEOUT_IN_SEC(devCtx->AutoBalloonContext->DelayTime);
    ntStatus = KeWaitForSingleObject(&devCtx->AutoBalloonContext->WakeUpMonitor,
        Executive, KernelMode, FALSE, &delay);

    delay.QuadPart = WDF_REL_TIMEOUT_IN_SEC(1);
    for (;;) {
        ntStatus = KeWaitForSingleObject(&devCtx->AutoBalloonContext->WakeUpMonitor, Executive,
            KernelMode, FALSE, &delay);
        if (devCtx->bShutDown) {
            break;
        }
        QueryRegistry(Device);
        if (devCtx->AutoBalloonContext->IsReady != VM_STATE_READY) {
            if (devCtx->num_pages != 0) {
                devCtx->AutoBalloonContext->Actual->val = 0;
                devCtx->AutoBalloonContext->bCtxInitialized = FALSE;
                BalloonTellMessage(Device);
            }
            continue;
        }
        // 1. Collect Stats
        currStats.ActualPages = devCtx->num_pages;
        if (!GatherAutoBalloonStats(&currStats)) {
            continue;
        }
        if (!devCtx->AutoBalloonContext->bCtxInitialized) {
            BalloonMemoryInit(Device, &currStats, &inflatePred, &deflatePred);
            balloonPages = AutoBalloonPredict(Device, &inflatePred, &currStats);
            if (balloonPages > BALLOON_INFLATE_THRESHOLD + devCtx->num_pages) {
                devCtx->AutoBalloonContext->Actual->val = min(devCtx->AutoBalloonContext->InflateMax, balloonPages);
                BalloonTellMessage(Device);
            }
            continue;
        }

        // 2. Deflate Period
        balloonPages = AutoBalloonPredict(Device, &deflatePred, &currStats);
        TraceEvents(TRACE_LEVEL_INFORMATION, DBG_UVP_DEBUG, "Actual %lu, Avai %lu balloon %lu\n",
            currStats.ActualPages, currStats.AvailablePages, balloonPages);
        inflatePred.CommittedPages = (inflatePred.CommittedPages + deflatePred.CommittedPages) / ARITHMETIC_MEAN;
        inflatePred.InUsedPages = (inflatePred.InUsedPages + deflatePred.InUsedPages) / ARITHMETIC_MEAN;
        if (balloonPages < devCtx->num_pages) {
            inflatePred.LocalTime = deflatePred.LocalTime;
            inflatePred.PageReadCount = deflatePred.PageReadCount;
            if ((balloonPages == 0) || ((devCtx->num_pages - balloonPages) > BALLOON_DEFLATE_THRESHOLD)) {
                devCtx->AutoBalloonContext->Actual->val = balloonPages;
                BalloonTellMessage(Device);
                AvailablePagesPeak = min(AvailablePagesPeak, currStats.AvailablePages);
                CommittedPagesPeak = max(CommittedPagesPeak, currStats.CommittedPages);
            }
            continue;
        }

        // 3. Inflate Period
        if ((currStats.LocalTime.QuadPart - inflatePred.LocalTime.QuadPart) >=
            ((LONGLONG)devCtx->AutoBalloonContext->Interval * SECONDS_FROM_100NANOSECONDS)) {
            if ((WEIGHT_SCALE + devCtx->AutoBalloonContext->BufferWeight) * inflatePred.InUsedPages <
                (WEIGHT_SCALE - devCtx->AutoBalloonContext->InflateWeight) * devCtx->AutoBalloonContext->RamPages) {
                TryInflate(Device);
                continue;
            }
            if (AvailablePagesPeak < currStats.AvailablePages) {
                currStats.AvailablePages = (AvailablePagesPeak + currStats.AvailablePages) / ARITHMETIC_MEAN;
            }
            if (CommittedPagesPeak > currStats.CommittedPages) {
                currStats.CommittedPages = (CommittedPagesPeak + currStats.CommittedPages) / ARITHMETIC_MEAN;
            }
            AvailablePagesPeak = devCtx->AutoBalloonContext->RamPages;
            CommittedPagesPeak = 0;

            balloonPages = AutoBalloonPredict(Device, &inflatePred, &currStats);
            TraceEvents(TRACE_LEVEL_INFORMATION, DBG_UVP_DEBUG, "Actual %lu, Avai %lu balloon %lu\n",
                currStats.ActualPages, currStats.AvailablePages, balloonPages);
            if ((balloonPages > devCtx->num_pages) && (balloonPages - devCtx->num_pages) > BALLOON_INFLATE_THRESHOLD) {
                devCtx->AutoBalloonContext->Actual->val = balloonPages;
                BalloonTellMessage(Device);
            }
        } else {
            AvailablePagesPeak = min(AvailablePagesPeak, currStats.AvailablePages);
            CommittedPagesPeak = max(CommittedPagesPeak, currStats.CommittedPages);
        }
    }

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_UVP_DEBUG, "Auto-Balloon monitor thread about to exit...");
    PsTerminateSystemThread(STATUS_SUCCESS);
}

NTSTATUS BalloonCreateMonitorThread(IN WDFDEVICE  Device)
{
    PDEVICE_CONTEXT     devCtx = GetDeviceContext(Device);
    NTSTATUS            status = STATUS_SUCCESS;
    HANDLE              hThread = 0;
    OBJECT_ATTRIBUTES   oa;

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "--> %s", __FUNCTION__);

    if (NULL == devCtx->AutoBalloonContext->MonitorThread)
    {
        InitializeObjectAttributes(&oa, NULL,
            OBJ_KERNEL_HANDLE, NULL, NULL);

        status = PsCreateSystemThread(&hThread, THREAD_ALL_ACCESS, &oa, NULL, NULL,
            AutoBalloonWorker, Device);
        if (!NT_SUCCESS(status))
        {
            TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
                "failed to create worker thread status 0x%08x\n", status);
            return status;
        }
        ObReferenceObjectByHandle(hThread, THREAD_ALL_ACCESS, NULL,
            KernelMode, (PVOID*)&devCtx->AutoBalloonContext->MonitorThread, NULL);
        KeSetPriorityThread(devCtx->AutoBalloonContext->MonitorThread, LOW_REALTIME_PRIORITY);

        ZwClose(hThread);
    }

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "<-- %s", __FUNCTION__);
    return status;
}

NTSTATUS BalloonCloseMonitorThread(IN WDFDEVICE  Device)
{
    PDEVICE_CONTEXT     devCtx = GetDeviceContext(Device);
    NTSTATUS            status = STATUS_SUCCESS;

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "--> %s", __FUNCTION__);

    PAGED_CODE();

    if (NULL != devCtx->AutoBalloonContext && NULL != devCtx->AutoBalloonContext->MonitorThread)
    {
        KeSetEvent(&devCtx->AutoBalloonContext->WakeUpMonitor, EVENT_INCREMENT, FALSE);
        status = KeWaitForSingleObject(devCtx->AutoBalloonContext->MonitorThread,
            Executive, KernelMode, FALSE, NULL);
        if (!NT_SUCCESS(status))
        {
            TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP,
                "KeWaitForSingleObject didn't succeed status 0x%08x\n", status);
        }
        ObDereferenceObject(devCtx->AutoBalloonContext->MonitorThread);
        devCtx->AutoBalloonContext->MonitorThread = NULL;
    }

    TraceEvents(TRACE_LEVEL_INFORMATION, DBG_INIT, "<-- %s\n", __FUNCTION__);
    return status;
}
